import numpy as np
import onnxruntime as rt
from sklearn.preprocessing import StandardScaler
import pickle
import os

# 示例输入数据
sample_inputs = [
    [26.20, 91.70, 803.10],  # 示例1：高温、高湿、高辐照
    [21.10, 99.20, 322.23],  # 示例2：中温、中湿、中高辐照
    [11.50, 88.99, 312.16]   # 示例3：低温、高湿、低辐照
]

def load_or_train_scaler():
    """加载已保存的标准化器或重新训练"""
    scaler_path = 'weather_scaler.pkl'
    
    if os.path.exists(scaler_path):
        # 加载已保存的标准化器
        with open(scaler_path, 'rb') as f:
            scaler = pickle.load(f)
        print("已加载保存的标准化器")
    else:
        print("找不到保存的标准化器，您需要运行train_weather_svm.py来生成它")
        # 如果需要可以在这里添加代码，直接训练一个临时的标准化器
        # 但最好是先运行训练脚本
        raise FileNotFoundError("请先运行train_weather_svm.py生成模型和标准化器")
    
    return scaler

def predict_weather(input_data):
    """使用ONNX模型预测天气"""
    # 加载模型
    session = rt.InferenceSession('weather_pytorch_model.onnx')
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    
    # 加载标准化器
    scaler = load_or_train_scaler()
    
    # 标准化输入数据
    input_scaled = scaler.transform(np.array(input_data, dtype=np.float32))
    
    # 预测
    results = []
    weather_types = ['晴天', '多云', '雨天']
    
    for i, sample in enumerate(input_scaled):
        # 扩展维度以匹配模型输入形状
        sample_reshaped = sample.reshape(1, -1).astype(np.float32)
        
        # 运行模型
        output = session.run([output_name], {input_name: sample_reshaped})[0]
        
        # 获取预测结果
        probabilities = output[0]
        prediction = np.argmax(probabilities)
        
        # 保存结果
        result = {
            'input': input_data[i],
            'probabilities': probabilities,
            'prediction': prediction,
            'weather_type': weather_types[prediction]
        }
        results.append(result)
    
    return results

# 添加一个程序保存标准化器的功能
def save_scaler_from_training():
    """运行训练并保存标准化器"""
    from train_weather_svm import generate_weather_data
    from sklearn.model_selection import train_test_split
    
    # 生成数据
    X, y = generate_weather_data(1000)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 训练标准化器
    scaler = StandardScaler()
    scaler.fit(X_train)
    
    # 保存标准化器
    with open('weather_scaler.pkl', 'wb') as f:
        pickle.dump(scaler, f)
    
    print("标准化器已保存至 weather_scaler.pkl")

# 运行示例
if __name__ == "__main__":
    # 检查是否存在scaler文件，不存在则尝试在train_weather_svm.py中保存
    if not os.path.exists('weather_scaler.pkl'):
        try:
            save_scaler_from_training()
        except:
            print("无法自动创建标准化器，请先运行train_weather_svm.py")
    
    # 检查是否存在模型文件
    if not os.path.exists('weather_pytorch_model.onnx'):
        print("找不到模型文件，请先运行train_weather_svm.py生成模型")
    else:
        # 进行天气预测
        results = predict_weather(sample_inputs)
        
        # 打印结果
        print("\n===== 天气预测示例 =====\n")
        for i, result in enumerate(results):
            print(f"示例 {i+1}:")
            print(f"输入数据: 温度={result['input'][0]:.2f}℃, 湿度={result['input'][1]:.2f}%, 辐照度={result['input'][2]:.2f}W/m²")
            print(f"预测天气: {result['weather_type']}")
            print(f"预测概率: 晴天={result['probabilities'][0]:.4f}, 多云={result['probabilities'][1]:.4f}, 雨天={result['probabilities'][2]:.4f}")
            print("-" * 50)
